import ts from "dedent";
import * as Path from "pathe";
import * as Pathe from "pathe/utils";

import * as Babel from "../vite/babel";
import type { Context } from "./context";
import * as Params from "./params";
import * as Route from "./route";
import type { RouteManifestEntry } from "../config/routes";

export type VirtualFile = { filename: string; content: string };

export function typesDirectory(ctx: Context) {
  return Path.join(ctx.rootDirectory, ".react-router/types");
}

export function generateFuture(ctx: Context): VirtualFile {
  const filename = Path.join(typesDirectory(ctx), "+future.ts");
  const content = ts`
    // Generated by React Router

    import "react-router";

    declare module "react-router" {
      interface Future {
        v8_middleware: ${ctx.config.future.v8_middleware}
      }
    }
  `;
  return { filename, content };
}

export function generateServerBuild(ctx: Context): VirtualFile {
  const filename = Path.join(typesDirectory(ctx), "+server-build.d.ts");
  const content = ts`
    // Generated by React Router

    declare module "virtual:react-router/server-build" {
      import { ServerBuild } from "react-router";
      export const assets: ServerBuild["assets"];
      export const assetsBuildDirectory: ServerBuild["assetsBuildDirectory"];
      export const basename: ServerBuild["basename"];
      export const entry: ServerBuild["entry"];
      export const future: ServerBuild["future"];
      export const isSpaMode: ServerBuild["isSpaMode"];
      export const prerender: ServerBuild["prerender"];
      export const publicPath: ServerBuild["publicPath"];
      export const routeDiscovery: ServerBuild["routeDiscovery"];
      export const routes: ServerBuild["routes"];
      export const ssr: ServerBuild["ssr"];
      export const unstable_getCriticalCss: ServerBuild["unstable_getCriticalCss"];
    }
  `;
  return { filename, content };
}

const { t } = Babel;
export function generateRoutes(ctx: Context): Array<VirtualFile> {
  // precompute
  const fileToRoutes = new Map<string, Set<string>>();
  const lineages = new Map<string, Array<RouteManifestEntry>>();
  const allPages = new Set<string>();
  const routeToPages = new Map<string, Set<string>>();
  for (const route of Object.values(ctx.config.routes)) {
    // fileToRoutes
    let routeIds = fileToRoutes.get(route.file);
    if (!routeIds) {
      routeIds = new Set();
      fileToRoutes.set(route.file, routeIds);
    }
    routeIds.add(route.id);

    // lineages
    const lineage = Route.lineage(ctx.config.routes, route);
    lineages.set(route.id, lineage);

    // pages
    const fullpath = Route.fullpath(lineage);
    if (!fullpath) continue;

    const pages = expand(fullpath);
    pages.forEach((page) => allPages.add(page));

    // routePages
    lineage.forEach(({ id }) => {
      let routePages = routeToPages.get(id);
      if (!routePages) {
        routePages = new Set<string>();
        routeToPages.set(id, routePages);
      }
      pages.forEach((page) => routePages.add(page));
    });
  }

  // +routes.ts
  const routesTs: VirtualFile = {
    filename: Path.join(typesDirectory(ctx), "+routes.ts"),
    content:
      ts`
        // Generated by React Router

        import "react-router"

        declare module "react-router" {
          interface Register {
            pages: Pages
            routeFiles: RouteFiles
            routeModules: RouteModules
          }
        }
      ` +
      "\n\n" +
      Babel.generate(pagesType(allPages)).code +
      "\n\n" +
      Babel.generate(routeFilesType({ fileToRoutes, routeToPages })).code +
      "\n\n" +
      Babel.generate(routeModulesType(ctx)).code,
  };

  // **/+types/*.ts
  const allAnnotations: Array<VirtualFile> = Array.from(fileToRoutes.entries())
    .filter(([file]) => isInAppDirectory(ctx, file))
    .map(([file, routeIds]) =>
      getRouteAnnotations({ ctx, file, routeIds, lineages }),
    );

  return [routesTs, ...allAnnotations];
}

function pagesType(pages: Set<string>) {
  return t.tsTypeAliasDeclaration(
    t.identifier("Pages"),
    null,
    t.tsTypeLiteral(
      Array.from(pages).map((page) => {
        return t.tsPropertySignature(
          t.stringLiteral(page),
          t.tsTypeAnnotation(
            t.tsTypeLiteral([
              t.tsPropertySignature(
                t.identifier("params"),
                t.tsTypeAnnotation(paramsType(page)),
              ),
            ]),
          ),
        );
      }),
    ),
  );
}

function routeFilesType({
  fileToRoutes,
  routeToPages,
}: {
  fileToRoutes: Map<string, Set<string>>;
  routeToPages: Map<string, Set<string>>;
}) {
  return t.tsTypeAliasDeclaration(
    t.identifier("RouteFiles"),
    null,
    t.tsTypeLiteral(
      Array.from(fileToRoutes).map(([file, routeIds]) =>
        t.tsPropertySignature(
          t.stringLiteral(file),
          t.tsTypeAnnotation(
            t.tsUnionType(
              Array.from(routeIds).map((routeId) => {
                const pages = routeToPages.get(routeId) ?? new Set();
                return t.tsTypeLiteral([
                  t.tsPropertySignature(
                    t.identifier("id"),
                    t.tsTypeAnnotation(
                      t.tsLiteralType(t.stringLiteral(routeId)),
                    ),
                  ),
                  t.tsPropertySignature(
                    t.identifier("page"),
                    t.tsTypeAnnotation(
                      pages
                        ? t.tsUnionType(
                            Array.from(pages).map((page) =>
                              t.tsLiteralType(t.stringLiteral(page)),
                            ),
                          )
                        : t.tsNeverKeyword(),
                    ),
                  ),
                ]);
              }),
            ),
          ),
        ),
      ),
    ),
  );
}

function routeModulesType(ctx: Context) {
  return t.tsTypeAliasDeclaration(
    t.identifier("RouteModules"),
    null,
    t.tsTypeLiteral(
      Object.values(ctx.config.routes).map((route) =>
        t.tsPropertySignature(
          t.stringLiteral(route.id),
          t.tsTypeAnnotation(
            isInAppDirectory(ctx, route.file)
              ? t.tsTypeQuery(
                  t.tsImportType(
                    t.stringLiteral(
                      `./${Path.relative(ctx.rootDirectory, ctx.config.appDirectory)}/${route.file}`,
                    ),
                  ),
                )
              : t.tsUnknownKeyword(),
          ),
        ),
      ),
    ),
  );
}

function isInAppDirectory(ctx: Context, routeFile: string): boolean {
  const path = Path.resolve(ctx.config.appDirectory, routeFile);
  return path.startsWith(ctx.config.appDirectory);
}

function getRouteAnnotations({
  ctx,
  file,
  routeIds,
  lineages,
}: {
  ctx: Context;
  file: string;
  routeIds: Set<string>;
  lineages: Map<string, Array<RouteManifestEntry>>;
}) {
  const filename = Path.join(
    typesDirectory(ctx),
    Path.relative(ctx.rootDirectory, ctx.config.appDirectory),
    Path.dirname(file),
    "+types",
    Pathe.filename(file) + ".ts",
  );

  const matchesType = t.tsTypeAliasDeclaration(
    t.identifier("Matches"),
    null,
    t.tsUnionType(
      Array.from(routeIds).map((routeId) => {
        const lineage = lineages.get(routeId)!;
        return t.tsTupleType(
          lineage.map((route) =>
            t.tsTypeLiteral([
              t.tsPropertySignature(
                t.identifier("id"),
                t.tsTypeAnnotation(t.tsLiteralType(t.stringLiteral(route.id))),
              ),
              t.tsPropertySignature(
                t.identifier("module"),
                t.tsTypeAnnotation(
                  t.tsTypeQuery(
                    t.tsImportType(
                      t.stringLiteral(
                        relativeImportSource(
                          rootDirsPath(ctx, filename),
                          Path.resolve(ctx.config.appDirectory, route.file),
                        ),
                      ),
                    ),
                  ),
                ),
              ),
            ]),
          ),
        );
      }),
    ),
  );

  const routeImportSource = relativeImportSource(
    rootDirsPath(ctx, filename),
    Path.resolve(ctx.config.appDirectory, file),
  );

  const content =
    ts`
      // Generated by React Router

      import type { GetInfo, GetAnnotations } from "react-router/internal";

      type Module = typeof import("${routeImportSource}")

      type Info = GetInfo<{
        file: "${file}",
        module: Module
      }>
    ` +
    "\n\n" +
    Babel.generate(matchesType).code +
    "\n\n" +
    ts`
      type Annotations = GetAnnotations<Info & { module: Module, matches: Matches }, ${ctx.rsc}>;

      export namespace Route {
        // links
        export type LinkDescriptors = Annotations["LinkDescriptors"];
        export type LinksFunction = Annotations["LinksFunction"];

        // meta
        export type MetaArgs = Annotations["MetaArgs"];
        export type MetaDescriptors = Annotations["MetaDescriptors"];
        export type MetaFunction = Annotations["MetaFunction"];

        // headers
        export type HeadersArgs = Annotations["HeadersArgs"];
        export type HeadersFunction = Annotations["HeadersFunction"];

        // middleware
        export type MiddlewareFunction = Annotations["MiddlewareFunction"];

        // clientMiddleware
        export type ClientMiddlewareFunction = Annotations["ClientMiddlewareFunction"];

        // loader
        export type LoaderArgs = Annotations["LoaderArgs"];

        // clientLoader
        export type ClientLoaderArgs = Annotations["ClientLoaderArgs"];

        // action
        export type ActionArgs = Annotations["ActionArgs"];

        // clientAction
        export type ClientActionArgs = Annotations["ClientActionArgs"];

        // HydrateFallback
        export type HydrateFallbackProps = Annotations["HydrateFallbackProps"];

        // Component
        export type ComponentProps = Annotations["ComponentProps"];

        // ErrorBoundary
        export type ErrorBoundaryProps = Annotations["ErrorBoundaryProps"];
      }
    `;
  return { filename, content };
}

function relativeImportSource(from: string, to: string) {
  let path = Path.relative(Path.dirname(from), to);

  let extension = Path.extname(path);

  // no extension
  path = Path.join(Path.dirname(path), Pathe.filename(path));
  if (!path.startsWith("../")) path = "./" + path;

  // In typescript, we want to support "moduleResolution": "nodenext" as well as not having "allowImportingTsExtensions": true,
  // so we normalize all JS like files to `.js`, but allow other extensions such as `.mdx` and others that might be used as routes.
  if (!extension || /\.(js|ts)x?$/.test(extension)) {
    extension = ".js";
  }

  return path + extension;
}

function rootDirsPath(ctx: Context, typesPath: string): string {
  const rel = Path.relative(typesDirectory(ctx), typesPath);
  return Path.join(ctx.rootDirectory, rel);
}

function paramsType(path: string) {
  const params = Params.parse(path);
  return t.tsTypeLiteral(
    Object.entries(params).map(([param, isRequired]) => {
      const property = t.tsPropertySignature(
        t.stringLiteral(param),
        t.tsTypeAnnotation(t.tsStringKeyword()),
      );
      property.optional = !isRequired;
      return property;
    }),
  );
}

function expand(fullpath: string): Set<string> {
  function recurse(segments: Array<string>, index: number): Array<string> {
    if (index === segments.length) return [""];
    const segment = segments[index];

    const isOptional = segment.endsWith("?");
    const isDynamic = segment.startsWith(":");
    const required = segment.replace(/\?$/, "");

    const keep = !isOptional || isDynamic;
    const kept = isDynamic ? segment : required;

    const withoutSegment = recurse(segments, index + 1);
    const withSegment = withoutSegment.map((rest) => [kept, rest].join("/"));

    if (keep) return withSegment;
    return [...withoutSegment, ...withSegment];
  }

  const segments = fullpath.split("/");
  const expanded = new Set<string>();
  for (let result of recurse(segments, 0)) {
    if (result !== "/") result = result.replace(/\/$/, "");
    expanded.add(result);
  }
  return expanded;
}
